package jm.data.structure.tree;

import com.sun.org.apache.xerces.internal.impl.dv.xs.SchemaDVFactoryImpl;

import java.util.Comparator;

/**
 * @Description AVL树实现
 * @date 2022/4/14 9:17
 */
public class AVLTree<E> extends BBST<E>{
    public AVLTree(){
        this(null);
    }

    public AVLTree(Comparator<E> Comparator){
        super(Comparator);
    }

    private static class AVLNode<E> extends Node<E>{
        // 高度
        // 每次新增时的节点都是叶子节点，叶子节点的高度为1。
        int height = 1;

        public AVLNode(E e, Node<E> parent) {
            super(e, parent);
        }

        /**
         * 获取平衡因子
         * @return
         */
        private int balanceFactor(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            return leftHeight - rightHeight;
        }

        /**
         * 更新当前节点的高度
         */
        private void updateHeight(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            height = 1 + Math.max(leftHeight,rightHeight);
        }

        /**
         * 获取子节点中高度最高的
         * @return
         */
        private Node<E> tallerChild(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            if (rightHeight > leftHeight){
                return right;
            }
            if (rightHeight < leftHeight){
                return left;
            }
            // 如果高度相等就返回与当前节点同方向的节点
            return isLeftChild() ? left : right;
        }

    }

    @Override
    protected Node<E> createNode(E e, Node<E> parent) {
        return new AVLNode<E>(e,parent);
    }

    @Override
    protected void afterAdd(Node<E> node) {
        while (( node = node.parent ) != null){
            if (isBalance(node)){
                // 更新高度
                // 这里的node最起码就是新节点的父节点
                updateHeight(node);
            }else {
                // 恢复平衡
                // 这里是高度最低的不平衡节点
                reBalance(node);
                // 退出循环使树恢复平衡
                break;
            }
        }
    }

    @Override
    protected void afterRemove(Node<E> node) {
        while (( node = node.parent ) != null){
            if (isBalance(node)){
                // 更新高度
                // 这里的node最起码就是新节点的父节点
                updateHeight(node);
            }else {
                // 恢复平衡
                // 这里是高度最低的不平衡节点
                reBalancePro(node);
                // 恢复之后不能退出，因为可能会导致父节点也失衡
            }
        }
    }

    @Override
    protected void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {
        super.afterRotate(grand, parent, child);
        // 更新高度
        updateHeight(grand);
        updateHeight(parent);
    }

    @Override
    protected void rotate(Node<E> r, Node<E> a, Node<E> b, Node<E> c, Node<E> d, Node<E> e, Node<E> f, Node<E> g) {
        super.rotate(r, a, b, c, d, e, f, g);
        // 因为修改了 b 的左右子树，所以需要更新高度
        updateHeight(b);
        // 因为修改了 f 的左右子树，所以需要更新高度
        updateHeight(f);
        updateHeight(d);
    }

    /**
     * 判断是否平衡
     * @param node
     * @return
     */
    private boolean isBalance(Node<E> node){
       return Math.abs(((AVLNode<E>)node).balanceFactor()) <= 1;
    }

    /**
     * 恢复平衡
     * @param grand 最先（高度最低，最往下）失衡的节点
     */
    private void reBalance(Node<E> grand){
        // 获取node比较高的子节点
        Node<E> parent = ((AVLNode<E>)grand).tallerChild();
        Node<E> node = ((AVLNode<E>)parent).tallerChild();
        if (parent.isLeftChild()){ // L
            if (node.isLeftChild()){ // LL
                rotateRight(grand);
            }else { // LR
                rotateLeft(parent);
                rotateRight(grand);
            }
        }else { // R
            if (node.isLeftChild()){ // RL
                rotateRight(parent);
                rotateLeft(grand);
            }else { // RR
                rotateLeft(grand);
            }
        }
    }

    /**
     * 统一版恢复平衡
     * @param grand
     */
    protected void reBalancePro(Node<E> grand){
        Node<E> parent = ((AVLTree.AVLNode<E>)grand).tallerChild();
        Node<E> node = ((AVLTree.AVLNode<E>)parent).tallerChild();
        if (parent.isLeftChild()){ // L
            if (node.isLeftChild()){ // LL
                rotate(grand,
                        node.left, node, node.right,
                        parent,
                        parent.right, grand, grand.right);
            }else { // LR
                rotate(grand,
                        parent.left, parent, node.left,
                        node,
                        node.right, grand, grand.right);
            }
        }else { // R
            if (node.isLeftChild()){ // RL
                rotate(grand,
                        grand.left, grand, node.left,
                        node,
                        node.right, parent, parent.right);
            }else { // RR
                rotate(grand,
                        grand.left, grand, parent.left,
                        parent,
                        node.left, node, node.right);
            }
        }
    }

    /**
     * 更新节点高度
     * @param node
     */
    protected void updateHeight(Node<E> node){
        ((AVLTree.AVLNode<E>)node).updateHeight();
    }

}
